#include "gtest/gtest.h"

//#include "DataStructure/Common/include/Constant.h"
#include "TimeSteppers.h"

using namespace SG::Algebra; // Open the namespace
//using SG::DataStructure::Common::TOLERANCE;

/**
 * Consider the Single-DOF oscillator forced by harmonic excitation,
 * \f[
 *    m\ddot{u} + c \dot{u} + k u = F_{0}sinw\left ( \omega t \right )
 *    u\left ( 0 \right ) = u_{0}
 *    \dot{u}\left ( 0 \right ) = \dot{u}_{0}
 * \f]
 *
 * The analytical solution is given as,
 * \f[
 *    u = e^{-\zeta \omega _{n} t}\left ( \frac{\dot{u}_{0}+\zeta \omega _{n}u_{0}}{\omega _{d}}sin \omega _{d}t+ u_{0}cos \omega
 * _{d}t\right)
 *    + Xe^{-\zeta \omega _{n} t}\left ( \frac{\zeta \omega _{n}sin\alpha -\omega cos \alpha }{ \omega _{d}} sin\omega _{d}t
 *    + sin \alpha cos \omega_{d}t\right )+X sin\left ( \omega t -\alpha \right )
 *
 *    \omega _{n}=\sqrt{\frac{k}{m}}
 *
 *    \zeta =\frac{c}{2m\omega _{n}}
 *
 *    \omega _{d} =\omega _{n}\sqrt{1-\zeta ^{2}}
 *
 *    \gamma = \frac{\omega }{\omega _{n}}
 *
 *    X =\frac{F_{0}/k}{\sqrt{\left ( 1-\gamma ^2 \right )^2+\left ( 2\zeta \gamma \right )^2 }}
 *
 *    \alpha = arctan\frac{2\zeta \gamma }{1-\gamma ^2}
 *
 *  \f]
 *
 * @see 胡少伟. 结构振动理论及其应用. 中国建筑工业出版社.
 *
 * @author nene
 * @date November, 2024.
 */

/** Characterize the governing equation of one Single-DOF oscillator forced by harmonic excitation. */
class ForcedDynamicSystem : public SecondOrderODESystem
{
public:
    ForcedDynamicSystem()
    {
#if defined(SG_USE_MPI)
        M.resize (1, 1);
        C.resize (1, 1);
        K.resize (1, 1);

        const int row_start = M.row_start ();
        const int row_end   = M.row_end ();

        for (auto r = row_start; r < row_end; ++ r)
        {
            if (r == 0)
            {
                M.set (0, 0, m);
                C.set (0, 0, c);
                K.set (0, 0, c);
            }
        }

        M.assemble ();
        C.assemble ();
        K.assemble ();
#else
        typedef Eigen::Triplet<Real_t> Triplet;
        std::vector<Triplet> triplets;

        // --- M
        M.resize (1, 1);
        triplets.clear ();
        triplets.push_back (Triplet (0, 0, m));
        M.setFromTriplets (triplets.begin (), triplets.end ());

        // --- C
        C.resize (1, 1);
        triplets.clear ();
        triplets.push_back (Triplet (0, 0, c));
        C.setFromTriplets (triplets.begin (), triplets.end ());

        // --- K
        K.resize (1, 1);
        triplets.clear ();
        triplets.push_back (Triplet (0, 0, k));
        K.setFromTriplets (triplets.begin (), triplets.end ());
#endif
    }
    // // The number of equations in the system
    // int GetEquationCount () const override
    // {
    //     return 1;
    // }

    /// The value of F(t, u, v, a) at the time t.
    void ComputeImplicitValue (const Real_t& t, const Vec_t& u, const Vec_t& v, const Vec_t& a, Vec_t& f) override
    {}

    /// The value of G(t, u, v) at the time t.
    void ComputeExplicitValue (const Real_t& t, const Vec_t& u, const Vec_t& v, Vec_t& g) override
    {
#if defined(SG_USE_MPI)
        g.resize (1);

        const int row_start = g.row_start ();
        const int row_end   = g.row_end ();

        for (auto r = row_start; r < row_end; ++r)
        {
            if (r == 0)
            {
                g.set (0, F0 * sin (omega * t));
            }
        }

        g.assemble ();
#else
        g.resize (1);
        g[0] = F0 * sin (omega * t);

#endif
    }

    /// The jacobian of F(t, u, v, a) at the time @a t.
    void ComputeImplicitJacobian (const Real_t& t, const Vec_t& u, const Vec_t& v, const Vec_t& a, SpMat_t& dfdu, SpMat_t& dfdv, SpMat_t& dfda) override
    {
        dfdu = K;
        dfdv = C;
        dfda = M;
    }

    /// The jacobian of G(t, u, v) at the time @a t.
    void ComputeExplicitJacobian (const Real_t& t, const Vec_t& u, const Vec_t& v, SpMat_t& dgdu, SpMat_t& dgdv) override
    {

    }

    /**
     * Solve the a that satisfies the algebraic equation F(t, u, v, a) = G(t, u, v) at the time t.
     * For u" = G(t, u, v), it's the value of G(t, u, v) at the time t.
     */
    void ComputeDerivatives (const Real_t& t, const Vec_t& u, const Vec_t& v, Vec_t& a) override 
    {
#if defined(SG_USE_MPI)
        a.resize (1);

        const int row_start = a.row_start ();
        const int row_end   = a.row_end ();

        for (auto r = row_start; r < row_end; ++ r)
        {
            if (r == 0)
            {
                a.set (0, (F0 * sin (omega * t) - k * u[0] - c * v[0]) / m);
            }
        }

        a.assemble ();
#else
        a[0] = (F0 * sin (omega * t) - k * u[0] - c * v[0]) / m;
#endif
    }

public:
    Real_t m     = 1.0;
    Real_t c     = 1;
    Real_t k     = 1.0;
    Real_t F0    = 1.0;
    Real_t omega = 1.0;

    SpMat_t M, C, K;
};

class ODETest : public ::testing::Test
{
public:
    typedef Eigen::Triplet<Real_t> Triplet;

    static void SetUpTestSuite ()
    {}

    void ExactSolution (std::vector<double>& x)
    {
        Real_t m     = odes.m;
        Real_t c     = odes.c;
        Real_t k     = odes.k;
        Real_t F0    = odes.F0;
        Real_t omega = odes.omega;

        Real_t omega_n = sqrt (k / m);
        Real_t zeta    = c / (2.0 * m * omega_n);
        Real_t zeta2   = zeta * zeta;
        Real_t omega_d = omega_n * sqrt (1.0 - zeta2);
        Real_t gama    = omega / omega_n;
        Real_t gama2   = gama * gama;
        Real_t X       = F0 / k * sqrt ((1.0 - gama2) * (1.0 - gama2) + 4 * zeta2 * gama2);
        Real_t alpha   = atan (2.0 * zeta * gama / (1.0 - gama2));

        std::vector<double> analytical_solution;
        for (int i = 0; i <= steps; ++i)
        {
            Real_t t       = i * dt;
            Real_t theta   = omega * t;
            Real_t theta_n = omega_n * t;
            Real_t theta_d = omega_d * t;
            Real_t u;

            // term1
            Real_t fact1 = (v0 + zeta * omega_n * u0);
            fact1        = fact1 * sin (theta_d) / omega_d;
            fact1        = fact1 + u0 * cos (theta_d);
            u            = exp (-zeta * theta_n) * fact1;

            // term2
            Real_t fact2 = zeta * sin (alpha) - omega * cos (alpha);
            fact2        = fact2 * sin (theta_d) / omega_d;
            fact2        = fact2 + sin (alpha) * cos (theta_d);
            u            = u + X * exp (-zeta * theta_n) * fact2;

            // term3
            u = u + X * sin (theta - alpha);

            x.push_back (u);
        }
    }

    void CheckSolution (std::vector<double>& numerical, std::vector<double>& analytical)
    {
#ifdef SG_VERBOSE
        std::cout << "U: time,numerical,analytical = " << std::endl;
        for (int i = 0; i <= steps; ++i)
        {
            std::cout << i * dt << ", " << numerical[i] << ", " << analytical[i] << std::endl;
        }
#endif

        // testing the tolerance between the numerical/analytical solution
        for (int i = 0; i <= steps; ++i)
        {
            ASSERT_NEAR (numerical[i], analytical[i], tolerance);
        }
    }

protected:
    Real_t              tolerance = 0.1;
    ForcedDynamicSystem odes;
    int                 n     = 1;  // the number of equations in the system
    Real_t              u0    = 1.0;
    Real_t              v0    = 0.0;
    int                 steps = 10;
    Real_t              t0    = 0.0;                // start time
    Real_t              tf    = 1.0;                // final time
    Real_t              dt    = (tf - t0) / steps;  // time step
};

TEST_F (ODETest, CentralDifferenceTest)
{
    Vec_t U (n), V (n), A (n);

    // initial conditions
#if defined(SG_USE_MPI)
    const int row_start = U.row_start ();
    const int row_end   = U.row_end ();

// #   if defined(SG_VERBOSE)
//     std::cout << "row_start = " << row_start << std::endl;
//     std::cout << "row_end = " << row_end << std::endl;
// #   endif

    for (auto r = row_start; r < row_end; ++r)
    {
        if (r == 0)
        {
            U.set (0, u0);
            V.set (0, v0);
        }
    }

    U.assemble ();
    V.assemble ();
#else
    U[0] = u0;
    V[0] = v0;
#endif

    // ODE Solver
    CentralDifference solver;
    solver.Initialize (&odes);

    // numerical solution
    std::vector<double> numerical_solution;

    numerical_solution.push_back (u0);
    for (int i = 1; i <= steps; ++i)
    {
        Real_t t = i * dt;

        solver.SetTime (t);
        solver.SetTimeStep (dt);

        solver.Step (U, V, A);

#if defined(SG_USE_MPI)
        for (auto r = row_start; r < row_end; ++ r)
        {
            numerical_solution.push_back (U[r]);
        }
#else
        numerical_solution.push_back (U[0]);
#endif
    }

    // analytical solution
    std::vector<double> analytical_solution;
    ExactSolution (analytical_solution);

#if defined(SG_USE_MPI)
    if (numerical_solution.size() == analytical_solution.size())
    {
        CheckSolution (numerical_solution, analytical_solution);
    }
#else
    CheckSolution (numerical_solution, analytical_solution);
#endif
}

TEST_F (ODETest, ModifiedCentralDifferenceTest)
{
    Vec_t U (n), V (n), A (n);

    // initial conditions
#if defined(SG_USE_MPI)
    const int row_start = U.row_start ();
    const int row_end   = U.row_end ();

    for (auto r = row_start; r < row_end; ++r)
    {
        if (r == 0)
        {
            U.set (0, u0);
            V.set (0, v0);
        }
    }

    U.assemble ();
    V.assemble ();
#else
    U[0] = u0;
    V[0] = v0;
#endif

    // ODE Solver
    ModifiedCentralDifference solver;
    solver.Initialize (&odes);

    // numerical solution
    std::vector<double> numerical_solution;

    numerical_solution.push_back (u0);
    for (int i = 1; i <= steps; ++i)
    {
        Real_t t = i * dt;

        solver.SetTime (t);
        solver.SetTimeStep (dt);

        solver.Step (U, V, A);

#if defined(SG_USE_MPI)
        for (auto r = row_start; r < row_end; ++r)
        {
            numerical_solution.push_back (U[r]);
        }
#else
        numerical_solution.push_back (U[0]);
#endif
    }

    // analytical solution
    std::vector<double> analytical_solution;
    ExactSolution (analytical_solution);

#if defined(SG_USE_MPI)
    if (numerical_solution.size () == analytical_solution.size ())
    {
        CheckSolution (numerical_solution, analytical_solution);
    }
#else
    CheckSolution (numerical_solution, analytical_solution);
#endif
}

TEST_F (ODETest, NewmarkBetaTest)
{
    Vec_t U (n), V (n), A (n);

    // initial conditions
#if defined(SG_USE_MPI)
    const int row_start = U.row_start ();
    const int row_end   = U.row_end ();

    for (auto r = row_start; r < row_end; ++r)
    {
        if (r == 0)
        {
            U.set (0, u0);
            V.set (0, v0);
        }
    }

    U.assemble ();
    V.assemble ();
#else
    U[0] = u0;
    V[0] = v0;
#endif

    // ODE Solver
    NewmarkBeta solver;
    solver.Initialize (&odes);

    // numerical solution
    std::vector<double> numerical_solution;

    numerical_solution.push_back (u0);
    for (int i = 1; i <= steps; ++i)
    {
        Real_t t = i * dt;

        solver.SetTime (t);
        solver.SetTimeStep (dt);

        solver.Step (U, V, A);

#if defined(SG_USE_MPI)
        for (auto r = row_start; r < row_end; ++r)
        {
            numerical_solution.push_back (U[r]);
        }
#else
        numerical_solution.push_back (U[0]);
#endif
    }

    // analytical solution
    std::vector<double> analytical_solution;
    ExactSolution (analytical_solution);

#if defined(SG_USE_MPI)
    if (numerical_solution.size () == analytical_solution.size ())
    {
        CheckSolution (numerical_solution, analytical_solution);
    }
#else
    CheckSolution (numerical_solution, analytical_solution);
#endif
}

TEST_F (ODETest, WilsonThetaTest)
{
    Vec_t U (n), V (n), A (n);

    // initial conditions
#if defined(SG_USE_MPI)
    const int row_start = U.row_start ();
    const int row_end   = U.row_end ();

    for (auto r = row_start; r < row_end; ++r)
    {
        if (r == 0)
        {
            U.set (0, u0);
            V.set (0, v0);
        }
    }

    U.assemble ();
    V.assemble ();
#else
    U[0] = u0;
    V[0] = v0;
#endif

    // ODE Solver
    WilsonTheta solver;
    solver.Initialize (&odes);

    // numerical solution
    std::vector<double> numerical_solution;

    numerical_solution.push_back (u0);
    for (int i = 1; i <= steps; ++i)
    {
        Real_t t = i * dt;

        solver.SetTime (t);
        solver.SetTimeStep (dt);

        solver.Step (U, V, A);

#if defined(SG_USE_MPI)
        for (auto r = row_start; r < row_end; ++r)
        {
            numerical_solution.push_back (U[r]);
        }
#else
        numerical_solution.push_back (U[0]);
#endif
    }

    // analytical solution
    std::vector<double> analytical_solution;
    ExactSolution (analytical_solution);

#if defined(SG_USE_MPI)
    if (numerical_solution.size () == analytical_solution.size ())
    {
        CheckSolution (numerical_solution, analytical_solution);
    }
#else
    CheckSolution (numerical_solution, analytical_solution);
#endif
}

TEST_F (ODETest, GeneralizedAlphaTest)
{
    Vec_t U (n), V (n), A (n);

    // initial conditions
#if defined(SG_USE_MPI)
    const int row_start = U.row_start ();
    const int row_end   = U.row_end ();

    for (auto r = row_start; r < row_end; ++r)
    {
        if (r == 0)
        {
            U.set (0, u0);
            V.set (0, v0);
        }
    }

    U.assemble ();
    V.assemble ();
#else
    U[0] = u0;
    V[0] = v0;
#endif

    // ODE Solver
    GeneralizedAlpha solver;
    solver.Initialize (&odes);

    // numerical solution
    std::vector<double> numerical_solution;

    numerical_solution.push_back (u0);
    for (int i = 1; i <= steps; ++i)
    {
        Real_t t = i * dt;

        solver.SetTime (t);
        solver.SetTimeStep (dt);

        solver.Step (U, V, A);

#if defined(SG_USE_MPI)
        for (auto r = row_start; r < row_end; ++r)
        {
            numerical_solution.push_back (U[r]);
        }
#else
        numerical_solution.push_back (U[0]);
#endif
    }

    // analytical solution
    std::vector<double> analytical_solution;
    ExactSolution (analytical_solution);

#if defined(SG_USE_MPI)
    if (numerical_solution.size () == analytical_solution.size ())
    {
        CheckSolution (numerical_solution, analytical_solution);
    }
#else
    CheckSolution (numerical_solution, analytical_solution);
#endif
}

TEST_F (ODETest, BatheTest)
{
    Vec_t U (n), V (n), A (n);

    // initial conditions
#if defined(SG_USE_MPI)
    const int row_start = U.row_start ();
    const int row_end   = U.row_end ();

    for (auto r = row_start; r < row_end; ++r)
    {
        if (r == 0)
        {
            U.set (0, u0);
            V.set (0, v0);
        }
    }

    U.assemble ();
    V.assemble ();
#else
    U[0] = u0;
    V[0] = v0;
#endif

    // ODE Solver
    Bathe solver;
    solver.Initialize (&odes);

    // numerical solution
    std::vector<double> numerical_solution;

    numerical_solution.push_back (u0);
    for (int i = 1; i <= steps; ++i)
    {
        Real_t t = i * dt;

        solver.SetTime (t);
        solver.SetTimeStep (dt);

        solver.Step (U, V, A);

#if defined(SG_USE_MPI)
        for (auto r = row_start; r < row_end; ++r)
        {
            numerical_solution.push_back (U[r]);
        }
#else
        numerical_solution.push_back (U[0]);
#endif
    }

    // analytical solution
    std::vector<double> analytical_solution;
    ExactSolution (analytical_solution);

#if defined(SG_USE_MPI)
    if (numerical_solution.size () == analytical_solution.size ())
    {
        CheckSolution (numerical_solution, analytical_solution);
    }
#else
    CheckSolution (numerical_solution, analytical_solution);
#endif
}